/**
 * @file   FDmethod.h
 * @author WinterMelonCDY <wintermeloncdy@wintermeloncdy-VirtualBox>
 * @date   Sat Mar 25 14:58:14 2023
 * 
 * @brief  
 * 
 * 
 */

#include <iostream>
#include <cmath>
#include <algorithm>
#include <vector>
#include <eigen3/Eigen/Dense>
#include <eigen3/Eigen/Sparse>
#include <eigen3/Eigen/Core>

/*
    定义函数类
*/
class Function
{
public:
    virtual double operator()(double _x, double _y) = 0; // 返回u在(_x,_y)处的值
    virtual double diff_x(double _x, double _y) // 返回u在(_x,_y)处关于x的偏导数
    {
        return 0; 
    }
    virtual double diff_y(double _x, double _y) // 返回u在(_x,_y)处关于y的偏导数
    {
        return 0; 
    }
    virtual double Laplace(double _x, double _y) // 返回函数u在(_x,_y)处负Laplace算子作用后的值
    {
        return 0; 
    }
};

/*
    定义FD类，这里使用单位正方形上的FD方法
*/
class FD
{
public:
    virtual void solve() = 0;  // FD算法
    virtual double result(double _x, double _y) //返回计算得到的格点(_x,_y)处的值
    {
        return 0.0;  
    }  
    virtual std::vector<double> norm_error() //返回误差的1、2以及无穷范数
    {
        return std::vector<double>();  
    }  
};

/*
    利用FD方法对(0,1)*(0,1)上的微分方程-laplace u = f进行求解，其中f以及边界条件给定
*/
class FD_domain1 : public FD
{
private:
    Function &u;
    double h;
    int condition, N; // N: 每一行或者列上的网格点数（含边界）
    Eigen::VectorXd U; // U：计算得到的解
public:
    /*
        u:真实函数值，用于计算f以及边界条件，并用于和计算得到的值进行比较误差分析
        h:网格宽度
        condition：边界条件（1为D条件, 2为N条件, 3为M条件）
    */
    FD_domain1(Function &_u, double _h, int _condition): u(_u), h(_h), condition(_condition)
    {
        if ((condition != 1) && (condition != 2) && (condition != 3))
	    {
	      std::cerr<< "The condition is invalid!" <<std::endl; //边界条件无效
		    exit(-1);
	    }

        N = static_cast<int>(1.0 / h + 1.0); 
    }

    // 求解 AU = F
    void solve()
    {
        if (condition == 1) // D条件
        {
            Eigen::SparseMatrix<double> A(N*N, N*N);
            std::vector<Eigen::Triplet<double> > tripletlist; 
            Eigen::MatrixXd F(N*N,1); 
            F = Eigen::MatrixXd::Zero(N*N, 1); 

            // 沿x轴正方向、y轴正方向依次从0开始标号至N*N-1
            for (int i = 0; i < N*N; i++)
		    {
                // 计算序号i对应的坐标
                double x_i = (i % N)*h;
                double y_i = (i / N)*h;

                if (i % N == 0 || i % N == N-1 || i / N == 0 || i / N == N-1)
                {
                    tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                    F(i, 0) = u(x_i, y_i);
                }
                else
                {
                    tripletlist.push_back(Eigen::Triplet<double>(i, i, 4.0/(h*h)));
                    tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
                    tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
                    tripletlist.push_back(Eigen::Triplet<double>(i, i-N, -1.0/(h*h)));
                    tripletlist.push_back(Eigen::Triplet<double>(i, i+N, -1.0/(h*h)));
                    F(i, 0) = u.Laplace(x_i, y_i);
                }
		    }

            A.setFromTriplets(tripletlist.begin(), tripletlist.end());
            A.makeCompressed();
            Eigen::SparseLU<Eigen::SparseMatrix<double> > Solver_sparse;
            Solver_sparse.compute(A);
            U = Solver_sparse.solve(F);
        }

        if (condition == 2) // N条件
        {
            Eigen::SparseMatrix<double> A(N*N, N*N); 
            std::vector<Eigen::Triplet<double> > tripletlist; 
            Eigen::MatrixXd F(N*N,1); 
            F = Eigen::MatrixXd::Zero(N*N, 1); 

            // 沿x轴正方向、y轴正方向依次从0开始标号至N*N-1
            for (int i = 0; i < N*N; i++)
		    {
                // 计算序号i对应的坐标
                double x_i = (i % N)*h;
                double y_i = (i / N)*h;

                if (i == 0 || i == N-1 || i == N*N-N || i == N*N-1)
                {
                    tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                    F(i, 0) = u(x_i, y_i);
                }
                else
                {
                    if (i % N == 0 || i % N == N-1 || i / N == 0 || i / N == N-1)
                    {
                        if (i % N == 0) // 左边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, -3.0/(2.0*h)));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i+1, 2.0/h));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i+2, -1.0/(2.0*h)));
                            F(i, 0) = u.diff_x(x_i, y_i);
                        }
                        if (i % N == N-1) // 右边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, 3.0/(2.0*h)));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -2.0/h));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-2, 1.0/(2.0*h)));
                            F(i, 0) = u.diff_x(x_i, y_i);
                        }
                        if (i / N == 0) // 下边界
                        {
                            if (i == 1)
                            {
                                // 碰到不适定的特殊情况，需要补充约束条件
                                tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                                F(i, 0) = u(h, 0.0);
                            }
                            else
                            {
                                tripletlist.push_back(Eigen::Triplet<double>(i, i, -3.0/(2.0*h)));
                                tripletlist.push_back(Eigen::Triplet<double>(i, i+N, 2.0/h));
                                tripletlist.push_back(Eigen::Triplet<double>(i, i+2*N, -1.0/(2.0*h)));
                                F(i, 0) = u.diff_y(x_i, y_i);
                            }
                        }
                        if (i / N == N-1) // 上边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, 3.0/(2.0*h)));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-N, -2.0/h));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-2*N, 1.0/(2.0*h)));
                            F(i, 0) = u.diff_y(x_i, y_i);
                        }
                    }
                    else
                    {
                        tripletlist.push_back(Eigen::Triplet<double>(i, i, 4.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i-N, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i+N, -1.0/(h*h)));
                        F(i, 0) = u.Laplace(x_i, y_i);
                    }
                }
		    }

            A.setFromTriplets(tripletlist.begin(), tripletlist.end());
            A.makeCompressed();
            Eigen::SparseLU<Eigen::SparseMatrix<double> > Solver_sparse;
            Solver_sparse.compute(A);
            U = Solver_sparse.solve(F);
        }

        if (condition == 3) // M条件，与x轴平行两边采用D条件，另两边N条件
        {
            Eigen::SparseMatrix<double> A(N*N, N*N); 
            std::vector<Eigen::Triplet<double> > tripletlist; 
            Eigen::MatrixXd F(N*N,1); 
            F = Eigen::MatrixXd::Zero(N*N, 1); 

            // 沿x轴正方向、y轴正方向依次从0开始标号至N*N-1
            for (int i = 0; i < N*N; i++)
		    {
                // 计算序号i对应的坐标
                double x_i = (i % N)*h;
                double y_i = (i / N)*h;

                if (i == 0 || i == N-1 || i == N*N-N || i == N*N-1)
                {
                    tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                    F(i, 0) = u(x_i, y_i);
                }
                else
                {
                    if (i % N == 0 || i % N == N-1 || i / N == 0 || i / N == N-1)
                    {
                        if (i % N == 0) // 左边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, -3.0/(2.0*h)));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i+1, 2.0/h));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i+2, -1.0/(2.0*h)));
                            F(i, 0) = u.diff_x(x_i, y_i);
                        }
                        if (i % N == N-1) // 右边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, 3.0/(2.0*h)));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -2.0/h));
                            tripletlist.push_back(Eigen::Triplet<double>(i, i-2, 1.0/(2.0*h)));
                            F(i, 0) = u.diff_x(x_i, y_i);
                        }
                        if (i / N == 0) // 下边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                            F(i, 0) = u(x_i, y_i);
                        }
                        if (i / N == N-1) // 上边界
                        {
                            tripletlist.push_back(Eigen::Triplet<double>(i, i, 1.0));
                            F(i, 0) = u(x_i, y_i);
                        }
                    }
                    else
                    {
                        tripletlist.push_back(Eigen::Triplet<double>(i, i, 4.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i-1, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i+1, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i-N, -1.0/(h*h)));
                        tripletlist.push_back(Eigen::Triplet<double>(i, i+N, -1.0/(h*h)));
                        F(i, 0) = u.Laplace(x_i, y_i);
                    }
                }
		    }

            A.setFromTriplets(tripletlist.begin(), tripletlist.end());
            A.makeCompressed();
            Eigen::SparseLU<Eigen::SparseMatrix<double> > Solver_sparse;
            Solver_sparse.compute(A);
            U = Solver_sparse.solve(F);
        }
    }

    double result(double _x, double _y) 
    {
        if (std::abs(_x/h - std::round(_x/h)) > 1e-9 || std::abs(_y/h - std::round(_y/h)) > 1e-9)
        {
	  std::cerr<< "The input is invalid!" <<std::endl; //不是格点的情况，输入无效
		    exit(-1);
        }

        return U[std::round(_y/h)*N+std::round(_x/h)];
    }  

    std::vector<double> norm_error() //返回误差的1、2以及无穷范数
    {
        Eigen::VectorXd U_hat(N*N); // U_hat存储函数真实值
        Eigen::VectorXd E(N*N);
  
        for (int i = 0; i < N*N; ++i)
        {
            // 计算序号i对应的坐标
            double x_i = (i % N)*h;
            double y_i = (i / N)*h;
            U_hat(i) = u(x_i, y_i);
        }
        // 计算一范数
        double norm1 = h*h*(U - U_hat).lpNorm<1>();
        // 计算二范数
        double norm2 = h*(U - U_hat).norm();
        // 计算无穷范数
        double normInf = (U - U_hat).lpNorm<Eigen::Infinity>();

        return {norm1, norm2, normInf};
    }  
};
